import pytesseract
from PIL import ImageGrab, Image, ImageEnhance, ImageFilter
import numpy as np
import requests
import json
import base64
import hashlib
import time
import hmac
from urllib.parse import urlencode
from config.settings import TESSERACT_PATH, load_config
from PyQt6.QtWidgets import QApplication
from PyQt6.QtGui import QGuiApplication, QImage

# 设置Tesseract路径
pytesseract.pytesseract.tesseract_cmd = TESSERACT_PATH

class CloudOCRConfigError(Exception):
    pass

class OCR:
    @staticmethod
    def _preprocess_image(image):
        """
        对图像进行预处理以提高OCR识别准确率
        
        Args:
            image: PIL Image对象
            
        Returns:
            PIL Image: 处理后的图像
        """
        # 1. 转换为RGB模式（如果不是的话）
        if image.mode != 'RGB':
            image = image.convert('RGB')
            
        # 2. 放大图像（提高分辨率）
        width, height = image.size
        image = image.resize((width * 2, height * 2), Image.Resampling.LANCZOS)
        
        # 3. 增强对比度
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(2.0)
        
        # 4. 增强锐度
        enhancer = ImageEnhance.Sharpness(image)
        image = enhancer.enhance(2.0)
        
        # 5. 应用高斯模糊去噪
        image = image.filter(ImageFilter.GaussianBlur(radius=0.5))
        
        # 6. 再次锐化以恢复细节
        image = image.filter(ImageFilter.SHARPEN)
        
        # 7. 转换为灰度图
        image = image.convert('L')
        
        # 8. 自适应二值化
        image = image.point(lambda x: 0 if x < 128 else 255, '1')
        
        return image

    @staticmethod
    def _fix_ocr_text(text):
        """
        修复OCR识别结果中的常见错误
        
        Args:
            text: OCR识别出的文本
            
        Returns:
            str: 修复后的文本
        """
        import re
        
        # 修复将I识别成丨或|的问题
        # 1. 处理单词中的丨或|
        text = re.sub(r'(?<![0-9a-zA-Z])[丨|](?![0-9a-zA-Z])', 'I', text)
        
        # 2. 处理单词开头的丨或|
        text = re.sub(r'^[丨|]', 'I', text)
        
        # 3. 处理单词结尾的丨或|
        text = re.sub(r'[丨|]$', 'I', text)
        
        # 4. 处理单词中间的丨或|
        text = re.sub(r'(?<=[a-zA-Z])[丨|](?=[a-zA-Z])', 'I', text)
        
        # 5. 处理单独的丨或|
        text = re.sub(r'\b[丨|]\b', 'I', text)
        
        return text

    @staticmethod
    def qt_rect_to_physical(rect):
        screen = QApplication.primaryScreen()
        dpi_ratio = screen.devicePixelRatio()
        return (
            int(rect.x() * dpi_ratio),
            int(rect.y() * dpi_ratio),
            int((rect.x() + rect.width()) * dpi_ratio),
            int((rect.y() + rect.height()) * dpi_ratio)
        )

    @staticmethod
    def grab_region_as_pil(physical_rect):
        """
        用PyQt的QScreen.grabWindow截图，返回PIL Image，支持多屏+DPI
        physical_rect: QRect(x, y, w, h)，全局物理像素坐标
        """
        screens = QGuiApplication.screens()
        center = physical_rect.center()
        for screen in screens:
            if screen.geometry().contains(center):
                target_screen = screen
                break
        else:
            target_screen = QGuiApplication.primaryScreen()
        screen_geo = target_screen.geometry()
        # 物理像素坐标转为该屏幕的相对物理像素坐标
        x = physical_rect.x() - screen_geo.x()
        y = physical_rect.y() - screen_geo.y()
        w = physical_rect.width()
        h = physical_rect.height()
        qpix = target_screen.grabWindow(0, x, y, w, h)
        qimg = qpix.toImage().convertToFormat(QImage.Format.Format_RGBA8888)
        ptr = qimg.bits()
        ptr.setsize(qimg.width() * qimg.height() * 4)
        img = Image.frombuffer('RGBA', (qimg.width(), qimg.height()), bytes(ptr), 'raw', 'RGBA', 0, 1)
        return img.convert('RGB')

    @staticmethod
    def _local_ocr(region, language="英语"):
        """本地OCR识别
        region: 应为ScreenSelector.get_selection_rect()返回的全局逻辑坐标QRect
        """
        # region应为ScreenSelector.get_selection_rect()，需先转物理像素坐标
        from PyQt6.QtCore import QRect
        if hasattr(region, 'get_physical_rect'):
            physical_rect = region.get_physical_rect()
        else:
            # 兼容旧调用，假设region已是QRect
            physical_rect = region
        img = OCR.grab_region_as_pil(physical_rect)
        # 图像预处理
        processed_image = OCR._preprocess_image(img)
        # 根据选择的语言设置OCR语言
        lang = 'eng' if language == "英语" else 'jpn'
        # OCR识别文字
        text = pytesseract.image_to_string(
            processed_image,
            lang=lang,
            config='--psm 6 --oem 3'  # 使用更精确的OCR模式
        )
        # 修复OCR识别结果中的常见错误
        text = OCR._fix_ocr_text(text)
        return text.strip()  # 清理文本前后的空白字符

    @staticmethod
    def _tencent_ocr(region, language="英语"):
        """腾讯OCR识别"""
        config = load_config()
        ocr_config = config.get('ocr_services', {}).get('tencent', {})
        
        secret_id = ocr_config.get('secret_id', '')
        secret_key = ocr_config.get('secret_key', '')
        
        if not secret_id or not secret_key:
            raise CloudOCRConfigError("腾讯OCR配置不完整，请在配置中设置SecretId和SecretKey")
        
        # region应为ScreenSelector.get_selection_rect()，需先转物理像素坐标
        if hasattr(region, 'get_physical_rect'):
            physical_rect = region.get_physical_rect()
        else:
            physical_rect = region
        img = OCR.grab_region_as_pil(physical_rect)
        
        # 将图像转换为base64
        import io
        img_buffer = io.BytesIO()
        img.save(img_buffer, format='PNG')
        img_base64 = base64.b64encode(img_buffer.getvalue()).decode('utf-8')
        
        # 根据语言设置OCR语言参数
        # 腾讯OCR支持的语言：auto（自动识别）、en（英文）、jap（日文）、kor(韩文)等
        language_mapping = {
            "英语": "en",
            "日语": "jap",
            "韩语": "kor",
        }
        ocr_language = language_mapping.get(language, "auto")
        
        # 构建请求参数
        service = "ocr"
        host = "ocr.tencentcloudapi.com"
        region = "ap-beijing"
        action = "GeneralBasicOCR"
        version = "2018-11-19"
        
        # 构建请求体
        payload = {
            "ImageBase64": img_base64,
            "LanguageType": ocr_language
        }
        
        # 生成签名
        timestamp = int(time.time())
        date = time.strftime("%Y-%m-%d", time.gmtime(timestamp))
        
        # 构建规范请求串
        http_request_method = "POST"
        canonical_uri = "/"
        canonical_querystring = ""
        canonical_headers = "content-type:application/json; charset=utf-8\nhost:" + host + "\n"
        signed_headers = "content-type;host"
        hashed_request_payload = hashlib.sha256(json.dumps(payload).encode("utf-8")).hexdigest()
        canonical_request = http_request_method + "\n" + canonical_uri + "\n" + canonical_querystring + "\n" + canonical_headers + "\n" + signed_headers + "\n" + hashed_request_payload
        
        # 构建待签名字符串
        algorithm = "TC3-HMAC-SHA256"
        request_timestamp = str(timestamp)
        credential_scope = date + "/" + service + "/" + "tc3_request"
        hashed_canonical_request = hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()
        string_to_sign = algorithm + "\n" + request_timestamp + "\n" + credential_scope + "\n" + hashed_canonical_request
        
        # 计算签名
        def sign(key, msg):
            return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()
        
        secret_date = sign(("TC3" + secret_key).encode("utf-8"), date)
        secret_service = sign(secret_date, service)
        secret_signing = sign(secret_service, "tc3_request")
        signature = hmac.new(secret_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()
        
        # 构建授权信息
        authorization = algorithm + " " + "Credential=" + secret_id + "/" + credential_scope + ", " + "SignedHeaders=" + signed_headers + ", " + "Signature=" + signature
        
        # 发送请求
        headers = {
            "Content-Type": "application/json; charset=utf-8",
            "Host": host,
            "Authorization": authorization,
            "X-TC-Action": action,
            "X-TC-Version": version,
            "X-TC-Timestamp": request_timestamp,
            "X-TC-Region": region
        }
        
        url = "https://" + host
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        
        if response.status_code == 200:
            result = response.json()
            print("[DEBUG] 腾讯OCR返回:", result)
            if 'Response' in result and 'TextDetections' in result['Response']:
                texts = [item['DetectedText'] for item in result['Response']['TextDetections']]
                return ' '.join(texts)
            elif 'Response' in result and 'Error' in result['Response']:
                error = result['Response']['Error']
                error_code = error.get('Code', '')
                error_message = error.get('Message', '')
                # 新增：未检测到文本时直接返回空字符串
                if error_code == 'FailedOperation.ImageNoText':
                    return ''
                # 其它异常才抛出
                if 'AuthFailure' in error_code:
                    if 'SecretIdNotFound' in error_code:
                        raise CloudOCRConfigError("腾讯OCR SecretId不正确或不存在，请检查配置")
                    elif 'SecretKeyError' in error_code:
                        raise CloudOCRConfigError("腾讯OCR SecretKey不正确，请检查配置")
                    else:
                        raise CloudOCRConfigError(f"腾讯OCR认证失败: {error_message}")
                elif 'UnauthorizedOperation' in error_code:
                    raise CloudOCRConfigError("腾讯OCR服务未开通或权限不足，请检查是否已开通OCR服务")
                else:
                    raise CloudOCRConfigError(f"腾讯OCR服务错误: {error_message}")
            else:
                raise CloudOCRConfigError(f"腾讯OCR识别失败: {result}")
        else:
            raise CloudOCRConfigError(f"腾讯OCR请求失败: HTTP {response.status_code}")

    @staticmethod
    def _baidu_ocr(region, language="英语"):
        """百度OCR识别"""
        config = load_config()
        ocr_config = config.get('ocr_services', {}).get('baidu', {})
        
        app_id = ocr_config.get('app_id', '')
        app_key = ocr_config.get('app_key', '')
        
        if not app_id or not app_key:
            raise CloudOCRConfigError("百度OCR配置不完整，请在配置中设置APP ID和密钥")
        
        # region应为ScreenSelector.get_selection_rect()，需先转物理像素坐标
        if hasattr(region, 'get_physical_rect'):
            physical_rect = region.get_physical_rect()
        else:
            physical_rect = region
        img = OCR.grab_region_as_pil(physical_rect)
        
        # 将图像转换为base64
        import io
        img_buffer = io.BytesIO()
        img.save(img_buffer, format='PNG')
        img_base64 = base64.b64encode(img_buffer.getvalue()).decode('utf-8')
        
        # 获取access token
        token_url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={app_id}&client_secret={app_key}"
        token_response = requests.post(token_url)
        
        if token_response.status_code != 200:
            error_info = token_response.json() if token_response.content else {}
            error_msg = error_info.get('error_description', '未知错误')
            if 'invalid client' in error_msg.lower():
                raise CloudOCRConfigError("百度OCR APP ID或密钥不正确，请检查配置")
            else:
                raise CloudOCRConfigError(f"获取百度OCR access token失败: {error_msg}")
        
        access_token_info = token_response.json()
        if 'error' in access_token_info:
            error_msg = access_token_info.get('error_description', '未知错误')
            if 'invalid client' in error_msg.lower():
                raise CloudOCRConfigError("百度OCR APP ID或密钥不正确，请检查配置")
            else:
                raise CloudOCRConfigError(f"百度OCR认证失败: {error_msg}")
        
        access_token = access_token_info.get('access_token')
        if not access_token:
            raise CloudOCRConfigError("百度OCR access token获取失败")
        
        # 根据语言选择不同的OCR API
        # 百度OCR支持：general_basic（通用文字识别）、accurate（高精度）、japan（日语）等
        language_mapping = {
            "英语": "general_basic",  # 通用文字识别，支持英文
            "日语": "general_basic",  # 百度OCR没有专门的japan API，使用general_basic
            "韩语": "general_basic"    # 新增：支持韩语
        }
        ocr_api = language_mapping.get(language, "general_basic")
        
        # 调用OCR API - 使用正确的API端点
        ocr_url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/{ocr_api}"
        
        headers = {
            'Content-Type': 'application/x-www-form-urlencoded'
        }
        
        # 构建请求参数
        data = {
            'access_token': access_token,
            'image': img_base64
        }
        
        # 对于日语和韩语识别，添加language_type参数
        if language == "日语":
            data['language_type'] = 'JAP'  # 百度OCR日语参数
        elif language == "韩语":
            data['language_type'] = 'KOR'  # 百度OCR韩语参数
        
        response = requests.post(ocr_url, headers=headers, data=data)
        
        if response.status_code == 200:
            result = response.json()
            print("[DEBUG] 百度OCR返回:", result)
            if 'words_result' in result:
                texts = [item['words'] for item in result['words_result']]
                return ' '.join(texts)
            elif 'error_code' in result:
                error_code = result.get('error_code')
                error_msg = result.get('error_msg', '未知错误')
                
                # 根据错误代码判断具体问题
                if error_code == 110:
                    raise CloudOCRConfigError("百度OCR access token无效，请检查APP ID和密钥配置")
                elif error_code == 111:
                    raise CloudOCRConfigError("百度OCR access token过期，请重新配置")
                elif error_code == 216100:
                    raise CloudOCRConfigError("百度OCR服务未开通，请先开通OCR服务")
                elif error_code == 216101:
                    raise CloudOCRConfigError("百度OCR服务调用量超限，请检查服务配额")
                elif error_code == 3:
                    raise CloudOCRConfigError("百度OCR API方法不支持，请检查API版本或服务开通状态")
                else:
                    raise CloudOCRConfigError(f"百度OCR服务错误 (错误码:{error_code}): {error_msg}")
            else:
                raise CloudOCRConfigError(f"百度OCR识别失败: {result}")
        else:
            raise CloudOCRConfigError(f"百度OCR请求失败: HTTP {response.status_code}")

    @staticmethod
    def recognize_text(region, language="英语", ocr_service="本地OCR"):
        """
        识别指定区域的文字
        
        Args:
            region: 要识别的区域 (x, y, width, height)
            language: 识别语言 ("英语" 或 "日语")
            ocr_service: OCR服务 ("本地OCR", "腾讯OCR", "百度OCR")
            
        Returns:
            str: 识别出的文字
        """
        try:
            if ocr_service == "本地OCR":
                return OCR._local_ocr(region, language)
            elif ocr_service == "腾讯OCR":
                return OCR._tencent_ocr(region, language)
            elif ocr_service == "百度OCR":
                return OCR._baidu_ocr(region, language)
            else:
                # 默认使用本地OCR
                return OCR._local_ocr(region, language)
        except CloudOCRConfigError as e:
            return ("__CLOUD_OCR_CONFIG_ERROR__", str(e))
        except Exception as e:
            print(f"OCR识别失败: {str(e)}")
            # 如果云OCR失败，回退到本地OCR
            if ocr_service != "本地OCR":
                print("云OCR失败，回退到本地OCR")
                return OCR._local_ocr(region, language)
            else:
                raise e 